import javax.crypto.*;
import javax.crypto.spec.*;
import java.security.*;
import java.lang.*;
import java.io.*;
import java.util.*;

public class DESOFB {

    public byte[] iv,key,input;
    public Cipher en;
    public Cipher dec;
    public File inputFile;
    public File outputFile;
    public FileInputStream in;
    public FileOutputStream out;
    public int pad;
    public int blocksize;

  public DESOFB(String mode,String args_iv, String args_key, int blocksize){

      try{
          System.out.println("Initializing key, iv.....");
          iv = args_iv.getBytes("UTF8");
          key = args_key.getBytes("UTF8");
          this.blocksize = blocksize / 8;

          ofbmode_start(mode);

      }catch(Exception e){System.out.println("Error encoding iv, key");}
      
  }

  public void ofbmode_start(String mode){

      if (mode.equals("enc") || mode.equals("dec")) {
          System.out.println("Starting Encryption using " + blocksize + " byte chunks");
          try{

              /******* OPEN FILES *********/
              if(mode.equals("enc")){
                  inputFile = new File("plain-in");
                  outputFile = new File("cipher-out");
              }
              else{
                  inputFile = new File("cipher-in");
                  outputFile = new File("plain-out");
              }
              in = new FileInputStream(inputFile);
              out = new FileOutputStream(outputFile);


              /****** PADDING *******/
              if(in.available() % blocksize != 0){
              
                  if (in.available() < blocksize) {
                      input = new byte[blocksize];
                  }
                  else{   
                      int temp = (in.available()) / blocksize;
                      //System.out.println(temp);
                      temp++;
                      pad = (temp*blocksize) - in.available();
                      input = new byte[in.available()+pad];
                  }
    
                  for (int i = in.available()-1;i<input.length;i++) {
                      input[i] = 0;
                  }
              }else{
                  input = new byte[in.available()];
              }
              in.read(input);
              System.out.println("# of bytes pad: " + input.length + " " + pad);


              /****** Cipher *******/

              en = Cipher.getInstance("DES/ECB/NOPadding");

              DESKeySpec dkey = new DESKeySpec(key);
              SecretKeyFactory fact = SecretKeyFactory.getInstance("DES");
              SecretKey deskey = fact.generateSecret(dkey);

              en.init(Cipher.ENCRYPT_MODE,deskey);



              /********** OFB MODE ************/

              byte[] tempin;
              byte[] tempout;
              byte[] tempIV = new byte[8];
              byte[] prevIV = new byte[8];
              int count = 0;

              //
              for (int i=0;i<input.length;i=i+blocksize) {
                  tempin = new byte[blocksize];

                  for (int j =0;j<tempin.length;j++) {
                      tempin[j] = input[i+j];
                      count++;
                  }
                  tempout = new byte[blocksize];
                  if(i==0){
                      tempIV = en.doFinal(iv);

                      
                      for(int k=0;k<blocksize;k++){
                          tempout[k] = tempIV[k];
                      }

                      prevIV = kShift(iv,tempout);
                      tempout = XOR(tempout,tempin);
                      
                  }
                  else{
                      tempIV = en.doFinal(prevIV);

                      for(int k=0;k<blocksize;k++){
                          tempout[k] = tempIV[k];
                      }

                      prevIV = kShift(prevIV,tempout);
                      tempout = XOR(tempout,tempin);
                      
                  }

                  if(mode.equals("dec")){
                      if(i == (input.length-blocksize)){
                          for(int k = 0;k<tempout.length;k++){
                              if(tempout[k] != new Byte("0").byteValue()){
                                  out.write(tempout[k]);
                              }
                          }
                      }else{
                          out.write(tempout);
                      }
                  }else{
                      out.write(tempout);
                  }
              }

              in.close();
              out.close();

          }catch(Exception e){System.out.println("Error during enc mode: " + e.toString());}
          

      }
      else{}
  }

  public byte[] XOR(byte[] arg1, byte[] arg2){
      byte[] result = new byte[blocksize];
      BitSet bits1 = new BitSet();
      BitSet bits2 = new BitSet();

      for (int i=0; i<arg1.length*8; i++) {
        if ((arg1[arg1.length-i/8-1]&(1<<(i%8))) > 0) {
            bits1.set(i);
        }
      }

      for (int i=0; i<arg1.length*8; i++) {
        if ((arg2[arg2.length-i/8-1]&(1<<(i%8))) > 0) {
            bits2.set(i);
        }
      }

      bits1.xor(bits2);
      for (int i=0; i<bits1.length(); i++) {
        if (bits1.get(i)) {
            result[result.length-i/8-1] |= 1<<(i%8);
        }
      }
      return result;
  }

  public byte[] kShift(byte[] oldIV, byte[] kBits){
      byte[] result = new byte[8];
      byte[] results = new byte[8];
      int counter=0;
      
      for(int i=0;i<result.length;i++){
          if(i<(8-blocksize)){
              result[i] = oldIV[i+blocksize];
          }
          else{
              result[i] = kBits[counter];
              counter++;
          }
      }

      return result;
  }



  public static void main (String args[]) {

     //Check number of arguments for CBC mode
    if (args.length != 4 || args[1].length() != 8 || args[2].length() != 8) {
      System.out.println("Invalid arguments");
      System.out.println("Usage: java DESCBC [enc/dec] <initialization vector> <secret key>");
      System.out.println("Note: vector and key much have length of 8");
    }
    else{

      // If encryption
      if(args[0].equals("enc") || args[0].equals("dec")){
        int bs = new Integer(args[3]).intValue();
        if(bs % 8 != 0){
            System.out.println("k-bits not a multiple of 8");
        }else{
            DESOFB DESenc = new DESOFB(args[0],args[1],args[2],bs);
        }
      }
         
      else{
        System.out.println("Mode must be: 'enc' or 'dec'");
      }
    }

  }

}
